import logging

from model.data_processing.jax_dataloader import NumpyLoader
from model.diffusion.diffusion_ctrl import DiffusionIKModel, DiffusionPlanningModel

log = logging.getLogger(__name__)


def train_diffusion_end_joint():
    """
    add joint weights !!
    """
    from model.data_processing.data_loaders import NormalizedEndJointData

    # val_data = NormalizedEndJointData("data/json_data/can0_revised_1752404975.json")
    # val_data.x = val_data.end[:, :3].numpy()
    # val_data.y = val_data.joints.numpy()

    data = NormalizedEndJointData("data/json_data/allover.json")

    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )

    self = DiffusionIKModel(in_dim=3, out_dim=6, dropout_rate=0.1, predict_epsilon=True, denoising_steps=10)
    self.train(data_loader, 20000, "diffusion_ej_EpsilonT10dropout.ckp", val_data=None)


def train_ik_diffusion():
    from model.data_processing.data_loaders import Batch, NormalizedEndJointData

    class NormalizedEJ6D(NormalizedEndJointData):
        """
        Normaized (x, y, z, rx, ry, rz) -> (j1:j6)
        """

        def __init__(self, **kwargs):
            super().__init__(**kwargs)

        def __getitem__(self, idx):
            return Batch(self.ends[idx], self.joints[idx])

    data = NormalizedEJ6D(filename="data/json_data/slow.json")

    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )

    self = DiffusionIKModel(in_dim=6, out_dim=6, dropout_rate=0.1, weight_decay=1e-4)
    # self.load_ckpt("diffusion_ik.ckp")
    self.train(data_loader, 20000, "diffusion_ik_slow.ckp", val_data=None)


def train_diffusion_end_joint_planning():
    """
    add joint weights !!
    """
    from model.data_processing.data_loaders import NormalizedEndJointSequenceData

    # val_data = NormalizedEndJointData("data/json_data/can0_revised_1752404975.json")
    # val_data.x = val_data.end[:, :3].numpy()
    # val_data.y = val_data.joints.numpy()
    horizon_steps = 1
    stride = 10  # x 0.01s

    data = NormalizedEndJointSequenceData(
        "data/json_data/can0_L20Full.json", horizon_steps=horizon_steps, stride=stride
    )

    data_loader = NumpyLoader(
        data,
        batch_size=512,
        shuffle=True,
        num_workers=10,
        pin_memory=False,
    )

    self = DiffusionPlanningModel(in_dim=6, out_dim=6, horizon_steps=horizon_steps)
    # self.load_ckpt("diffusion_arm_joint.ckp")
    self.train(data_loader, 20000, "diffusion_end_joint_planning_h1.ckp", val_data=None)


# def train_fourier_diffusion_end_joint():
#     from model import DiffusionFourierEndJointModel
#     from model.data_processing.data_loaders import FourierEndJointData

#     data = FourierEndJointData("data/json_data/can0_L20Full.json")

#     data_loader = NumpyLoader(
#         data,
#         batch_size=512,
#         shuffle=True,
#         num_workers=10,
#         pin_memory=False,
#     )

#     self = DiffusionFourierEndJointModel()
#     self.train(data_loader, 30000, "diffusion_end_joint_fourier.ckp")


if __name__ == "__main__":
    train_diffusion_end_joint()
    # train_ik_diffusion()
